import torch
from 数据集 import train_loader
from LeNet import LeNet

"""
模型训练的一般步骤
超参数
数据
损失函数
优化器
学习率调度器
开始循环训练
    清空梯度
    前向传播
    计算损失
    反向传播
    更新参数
保存模型

基本的提速方式：
1.使用GPU
2.利用多线程或多进程
3.调整批次大小，充分利用显卡
4.学习率调度器
"""

# ---创建设备---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ---超参数---
lr = 0.01
Epochs = 50

# ---创建模型---
model = LeNet()
model.to(device)  # 将模型移动到GPU

# ---创建损失函数、优化器、学习率调度器---
loss_fn = torch.nn.NLLLoss()  # 因为模型内部使用了softmax，所以损失函数使用NLLLoss，而不是CrossEntropyLoss
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)

# ---开始循环训练---
for epoch in range(Epochs):
    total_loss = 0
    for i, (imgs, labels) in enumerate(train_loader):
        # 将数据放到设备上
        imgs, labels = imgs.to(device), labels.to(device)
        optimizer.zero_grad()
        y_pre = model(imgs)
        loss = loss_fn(y_pre, labels)
        total_loss += loss.item()
        loss.backward()
        optimizer.step()
        if i == 0 or (i + 1) % 50 == 0:
            print(f'epoch: {epoch + 1}, step: {i + 1}, loss: {loss.item():.4f}')
    # 输出一轮训练的信息
    avg_loss = total_loss / len(train_loader)
    print(f'epoch: [{epoch + 1}/{Epochs}], loss: {avg_loss:.4f}')
    # 更新学习率
    scheduler.step()

# 保存模型
# torch.save(model.state_dict(), 'LeNet.pth')
